<?php

namespace Kunfeng\Dify\Model\Datasets;

/**
 * 语义检索请求模型
 */
class RetrieveRequest
{
    private $query;
    private $retrievalMode;
    private $topK;
    private $scoreThreshold;
    private $rerankingModel;

    public function __construct(
        string $query,
        string $retrievalMode = 'single',
        int $topK = 3,
        ?float $scoreThreshold = null,
        ?string $rerankingModel = null
    ) {
        $this->query = $query;
        $this->retrievalMode = $retrievalMode;
        $this->topK = $topK;
        $this->scoreThreshold = $scoreThreshold;
        $this->rerankingModel = $rerankingModel;
    }

    /**
     * 创建请求
     */
    public static function create(string $query): self
    {
        return new self($query);
    }

    /**
     * 转换为数组
     */
    public function toArray(): array
    {
        $data = [
            'query' => $this->query,
            'retrieval_mode' => $this->retrievalMode,
            'top_k' => $this->topK,
        ];

        if ($this->scoreThreshold !== null) {
            $data['score_threshold'] = $this->scoreThreshold;
        }
        if ($this->rerankingModel !== null) {
            $data['reranking_model'] = $this->rerankingModel;
        }

        return $data;
    }

    // Getters
    public function getQuery(): string
    {
        return $this->query;
    }

    public function getRetrievalMode(): string
    {
        return $this->retrievalMode;
    }

    public function getTopK(): int
    {
        return $this->topK;
    }

    public function getScoreThreshold(): ?float
    {
        return $this->scoreThreshold;
    }

    public function getRerankingModel(): ?string
    {
        return $this->rerankingModel;
    }
}
